use std::{path::Path, str::Utf8Error};

use thiserror::Error;

use crate::{
    ast::{
        BinaryExpr, BinaryOperator, CallExpr, Expr, ExprList, Function, LiteralExpr, Module,
        NumberExpr, PrintExpr, Prototype, ReturnExpr, VarDeclExpr, VarType, VariableExpr,
    },
    lexer::{HasLocation, LexError, Lexer, Location},
    token::Token,
};

#[derive(Debug, Error)]
pub enum ParseError {
    #[error("{0}")]
    LexError(#[from] LexError),
    #[error("expected: {0:?}, actual: {1:?}")]
    Unexpected(Vec<Token>, Token),
    #[error("No prefix parse function: {0:?}")]
    NoPrefixParseFn(Token),
    #[error("invalid dims")]
    InvalidDims,
    #[error("{0}")]
    Utf8Error(#[from] Utf8Error),
}

pub type Result<T> = std::result::Result<T, ParseError>;

pub struct Parser {
    lexer: Lexer,
    cur_tok: Token,
}
impl Parser {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
        let mut p = Self {
            lexer: Lexer::new(path)?,
            cur_tok: Token::Eof,
        };
        p.next_token();
        Ok(p)
    }

    fn next_token(&mut self) {
        self.cur_tok = self.lexer.next_token();
    }

    pub fn parse(&mut self) -> Result<Module> {
        self.parse_module()
    }

    fn parse_module(&mut self) -> Result<Module> {
        let mut functions = Vec::new();
        while self.cur_tok != Token::Eof {
            functions.push(self.parse_function()?);
        }
        Ok(Module { functions })
    }

    fn parse_function(&mut self) -> Result<Function> {
        let proto = self.parse_prototype()?;
        let body = self.parse_expr_list()?;
        Ok(Function { proto, body })
    }

    fn parse_prototype(&mut self) -> Result<Prototype> {
        let location = self.lexer.location_ref().clone();

        self.expect(Token::Def)?;
        let var = self.parse_variable_expr()?;
        self.expect(Token::LParen)?;
        let args = self.parse_list(
            |parser: &mut Parser| parser.parse_variable_expr(),
            Token::Comma,
            Token::RParen,
        )?;
        Ok(Prototype {
            name: var.name,
            args,
            location,
        })
    }

    fn expect(&mut self, expected: Token) -> Result<()> {
        if expected == self.cur_tok {
            self.next_token();
            Ok(())
        } else {
            Err(ParseError::Unexpected(vec![expected], self.cur_tok.clone()))
        }
    }

    fn parse_variable_expr(&mut self) -> Result<VariableExpr> {
        let location = self.lexer.location_ref().clone();

        let name = match &self.cur_tok {
            Token::Ident(ident) => ident.clone(),
            _ => {
                return Err(ParseError::Unexpected(
                    vec![Token::Ident(String::new())],
                    self.cur_tok.clone(),
                ));
            }
        };
        self.next_token();
        Ok(VariableExpr { name, location })
    }

    fn parse_expr_list(&mut self) -> Result<ExprList> {
        let mut exprs = Vec::new();
        self.expect(Token::LBrace)?;
        while self.cur_tok != Token::Eof && self.cur_tok != Token::RBrace {
            exprs.push(self.parse_expr(Precedence::Lowest)?);
            self.expect(Token::Semicolon)?;
        }
        self.expect(Token::RBrace)?;
        Ok(ExprList { exprs })
    }

    fn parse_expr(&mut self, precedence: Precedence) -> Result<Expr> {
        let prefix = get_prefix_parse_fns(&self.cur_tok)
            .ok_or(ParseError::NoPrefixParseFn(self.cur_tok.clone()))?;
        let mut left_expr = prefix(self)?;
        while self.cur_tok != Token::Semicolon
            && self.cur_tok != Token::Eof
            && precedence < get_precedence(&self.cur_tok)
        {
            if let Some(infix) = get_infix_parse_fns(&self.cur_tok) {
                left_expr = infix(self, left_expr)?;
            } else {
                return Ok(left_expr);
            }
        }
        Ok(left_expr)
    }

    fn parse_list<T>(
        &mut self,
        parse_fn: fn(&mut Parser) -> Result<T>,
        seperator: Token,
        terminator: Token,
    ) -> Result<Vec<T>> {
        let mut ret = Vec::new();
        if self.cur_tok == terminator {
            self.next_token();
            return Ok(ret);
        }
        ret.push(parse_fn(self)?);
        while self.cur_tok == seperator {
            self.next_token();
            ret.push(parse_fn(self)?);
        }
        self.expect(terminator)?;
        Ok(ret)
    }

    fn parse_number_expr(&mut self) -> Result<NumberExpr> {
        let location = self.lexer.location_ref().clone();

        let val = match &self.cur_tok {
            Token::Num(value) => *value,
            _ => {
                return Err(ParseError::Unexpected(
                    vec![Token::Num(0.0)],
                    self.cur_tok.clone(),
                ));
            }
        };
        self.next_token();
        Ok(NumberExpr { val, location })
    }

    fn parse_literal_expr(&mut self) -> Result<LiteralExpr> {
        let location = self.lexer.location_ref().clone();

        self.expect(Token::LBracket)?;
        let value_dims = self.parse_list(
            |parser: &mut Parser| {
                let expr = parser.parse_expr(Precedence::Lowest)?;
                let dims = parser.parse_dims(&expr)?;
                Ok((expr, dims))
            },
            Token::Comma,
            Token::RBracket,
        )?;

        if value_dims.is_empty() {
            Ok(LiteralExpr {
                values: Vec::new(),
                dims: Vec::new(),
                location,
            })
        } else {
            let mut dims = value_dims[0].1.clone();
            if value_dims.iter().find(|(_, d)| d != &dims).is_some() {
                return Err(ParseError::InvalidDims);
            }
            let values: Vec<Expr> = value_dims.into_iter().map(|(v, _)| v).collect();
            dims.insert(0, values.len() as u64);
            Ok(LiteralExpr {
                values,
                dims,
                location,
            })
        }
    }

    fn parse_var_decl_expr(&mut self) -> Result<VarDeclExpr> {
        let location = self.lexer.location_ref().clone();

        self.expect(Token::Var)?;
        let var = self.parse_variable_expr()?;
        let r#type = if self.cur_tok == Token::Lt {
            self.parse_var_type()?
        } else {
            VarType::default()
        };
        self.expect(Token::Assign)?;
        let init_val = Box::new(self.parse_expr(Precedence::Lowest)?);
        Ok(VarDeclExpr {
            name: var.name,
            r#type,
            init_val,
            location,
        })
    }

    fn parse_var_type(&mut self) -> Result<VarType> {
        self.expect(Token::Lt)?;
        let shape = self.parse_list(
            |parser: &mut Parser| parser.parse_u64(),
            Token::Comma,
            Token::Gt,
        )?;
        Ok(VarType { shape })
    }

    fn parse_u64(&mut self) -> Result<u64> {
        let number_expr = self.parse_number_expr()?;
        Ok(number_expr.val as u64)
    }

    fn parse_dims(&mut self, expr: &Expr) -> Result<Vec<u64>> {
        match expr {
            Expr::NumberExpr(_number_expr) => Ok(Vec::new()),
            Expr::LiteralExpr(literal_expr) => self.parse_dims_of(literal_expr),
            _ => todo!(),
        }
    }

    fn parse_dims_of(&mut self, literal_expr: &LiteralExpr) -> Result<Vec<u64>> {
        let mut dims = vec![literal_expr.values.len() as u64];
        if literal_expr.values.is_empty() {
            return Ok(dims);
        }
        let mut sub_dims = self.parse_dims(&literal_expr.values[0])?;
        for value in literal_expr.values.iter().skip(1) {
            let another_sub_dims = self.parse_dims(value)?;
            if sub_dims != another_sub_dims {
                return Err(ParseError::InvalidDims);
            }
        }
        dims.append(&mut sub_dims);
        Ok(dims)
    }

    fn parse_return_expr(&mut self) -> Result<ReturnExpr> {
        self.expect(Token::Return)?;
        let expr = if self.cur_tok != Token::Semicolon {
            Some(Box::new(self.parse_expr(Precedence::Lowest)?))
        } else {
            None
        };
        Ok(ReturnExpr { expr })
    }

    fn parse_binary_expr(&mut self, left: Expr) -> Result<BinaryExpr> {
        let op = match &self.cur_tok {
            Token::Plus => BinaryOperator::Add,
            Token::Asterisk => BinaryOperator::Mul,
            _ => todo!(),
        };
        let precedence = get_precedence(&self.cur_tok);
        self.next_token();
        let rhs = Box::new(self.parse_expr(precedence)?);
        Ok(BinaryExpr {
            lhs: Box::new(left),
            op,
            rhs,
        })
    }

    fn parse_print_expr(&mut self, location: Location) -> Result<PrintExpr> {
        self.expect(Token::LParen)?;
        let arg = Box::new(self.parse_expr(Precedence::Lowest)?);
        self.expect(Token::RParen)?;
        Ok(PrintExpr { arg, location })
    }

    fn parse_call_expr(&mut self, callee: String, location: Location) -> Result<CallExpr> {
        self.expect(Token::LParen)?;
        let args = self.parse_list(
            |parser: &mut Parser| parser.parse_expr(Precedence::Lowest),
            Token::Comma,
            Token::RParen,
        )?;
        Ok(CallExpr {
            callee,
            args,
            location,
        })
    }

    fn parse_grouped_expr(&mut self) -> Result<Expr> {
        self.expect(Token::LParen)?;
        let expr = self.parse_expr(Precedence::Lowest)?;
        self.expect(Token::RParen)?;
        Ok(expr)
    }
}

fn get_infix_parse_fns(tok: &Token) -> Option<fn(&mut Parser, Expr) -> Result<Expr>> {
    Some(match tok {
        Token::Plus | Token::Asterisk => {
            |parser: &mut Parser, left: Expr| Ok(Expr::BinaryExpr(parser.parse_binary_expr(left)?))
        }
        Token::LParen => |parser: &mut Parser, left: Expr| match &left {
            Expr::VariableExpr(VariableExpr { name, location: _ }) => {
                if name == "print" {
                    Ok(Expr::PrintExpr(parser.parse_print_expr(left.location())?))
                } else {
                    Ok(Expr::CallExpr(
                        parser.parse_call_expr(name.clone(), left.location())?,
                    ))
                }
            }
            _ => todo!(),
        },
        _ => return None,
    })
}

fn get_prefix_parse_fns(tok: &Token) -> Option<fn(&mut Parser) -> Result<Expr>> {
    Some(match tok {
        Token::Num(_num) => |parser: &mut Parser| Ok(Expr::NumberExpr(parser.parse_number_expr()?)),
        Token::LBracket => {
            |parser: &mut Parser| Ok(Expr::LiteralExpr(parser.parse_literal_expr()?))
        }
        Token::Var => |parser: &mut Parser| Ok(Expr::VarDeclExpr(parser.parse_var_decl_expr()?)),
        Token::Return => |parser: &mut Parser| Ok(Expr::ReturnExpr(parser.parse_return_expr()?)),
        Token::LParen => |parser: &mut Parser| Ok(parser.parse_grouped_expr()?),
        Token::Ident(_) => {
            |parser: &mut Parser| Ok(Expr::VariableExpr(parser.parse_variable_expr()?))
        }
        _ => return None,
    })
}

#[derive(PartialEq, PartialOrd)]
pub enum Precedence {
    Lowest,
    Sum,
    Product,
    Call,
}

fn get_precedence(tok: &Token) -> Precedence {
    match tok {
        Token::Asterisk => Precedence::Product,
        Token::Plus => Precedence::Sum,
        Token::LParen => Precedence::Call,
        _ => Precedence::Lowest,
    }
}
